package cn.iocoder.yudao.module.mqtt.framework.mqtt;

import cn.iocoder.yudao.module.mqtt.framework.util.MqttContextHolder;
import lombok.extern.slf4j.Slf4j;
import org.eclipse.paho.client.mqttv3.MqttMessage;
import org.springframework.integration.mqtt.support.DefaultPahoMessageConverter;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageHeaders;
import org.springframework.messaging.support.MessageBuilder;

import java.util.HashMap;
import java.util.Map;

/**
 * 自定义MQTT消息转换器
 * 在消息头中包含客户端ID信息
 *
 * @author 芋道源码
 */
@Slf4j
public class CustomMqttMessageConverter extends DefaultPahoMessageConverter {

    private final String clientId;
    private final Long configId;

    public CustomMqttMessageConverter(String clientId, Long configId) {
        this.clientId = clientId;
        this.configId = configId;
    }

    @Override
    public Message<?> toMessage(String topic, MqttMessage mqttMessage) {
        // 设置线程本地变量
        MqttContextHolder.setClientId(clientId);
        MqttContextHolder.setConfigId(configId);
        
        Message<?> originalMessage = super.toMessage(topic, mqttMessage);
        
        if (originalMessage != null) {
            // 获取原始消息头
            MessageHeaders originalHeaders = originalMessage.getHeaders();
            Map<String, Object> newHeaders = new HashMap<>(originalHeaders);
            
            // 添加客户端ID和配置ID到消息头
            newHeaders.put("mqtt_receivedClientId", clientId);
            newHeaders.put("mqtt_clientId", clientId);
            newHeaders.put("mqtt_configId", configId);
            newHeaders.put("mqtt_receivedConfigId", configId);
            
            log.debug("添加MQTT消息头 - 客户端ID: {}, 配置ID: {}", clientId, configId);
            
            // 创建新的消息
            return MessageBuilder.createMessage(
                    originalMessage.getPayload(),
                    new MessageHeaders(newHeaders)
            );
        }
        
        return originalMessage;
    }

    @Override
    public Message<?> toMessage(Object payload, MessageHeaders headers) {
        // 设置线程本地变量
        MqttContextHolder.setClientId(clientId);
        MqttContextHolder.setConfigId(configId);
        
        Message<?> originalMessage = super.toMessage(payload, headers);
        
        if (originalMessage != null) {
            // 获取原始消息头
            MessageHeaders originalHeaders = originalMessage.getHeaders();
            Map<String, Object> newHeaders = new HashMap<>(originalHeaders);
            
            // 添加客户端ID和配置ID到消息头
            newHeaders.put("mqtt_receivedClientId", clientId);
            newHeaders.put("mqtt_clientId", clientId);
            newHeaders.put("mqtt_configId", configId);
            newHeaders.put("mqtt_receivedConfigId", configId);
            
            log.debug("添加MQTT消息头 - 客户端ID: {}, 配置ID: {}", clientId, configId);
            
            // 创建新的消息
            return MessageBuilder.createMessage(
                    originalMessage.getPayload(),
                    new MessageHeaders(newHeaders)
            );
        }
        
        return originalMessage;
    }
} 